#include <Classifier/classifier.h>
#include <Log/output.h>


using namespace cv;
Classifier::Classifier(const string& modelfile, int maxThreadNum): model(modelfile), threadpool(maxThreadNum), maxThreadNum(maxThreadNum)
{}

void Classifier::set_show(bool flag)
{
    is_show = flag;
}

void getMaxConnectedComponent(const Mat& src, Mat& dst)
{
    Mat labels, centroids, stats;
    int nccomps = connectedComponentsWithStats(src, labels, stats, centroids);

    double max_area = -1;
    int max_label = -1;

    for (int i = 1; i < nccomps; ++i)
    {
        double area = stats.at<int>(i, CC_STAT_WIDTH) * stats.at<int>(i, CC_STAT_HEIGHT);
        if (area > max_area)
        {
            max_area = area;
            max_label = i;
        }
    }
    //cout << "max_area: " << max_area << endl;

    dst = Mat::zeros(src.size(), CV_8UC1);
    int* plabels = (int*)labels.data;
    uchar *pdst = dst.data;
    for (int i = 0, total = dst.rows * dst.cols; i < total; ++i,++plabels,++pdst)
    {
        if (*plabels == max_label) *pdst = 255;
        //else *pdst = 0;
    }

   /* for (int row = 0, rows = dst.rows; row < rows;++row)
    {
        int tmp = row * dst.cols;
        for (int col = 0, cols= dst.cols; col < cols; ++col)
        {
            if (*(plabels+ tmp + col) == max_label)
            {
                *(pdst + tmp + col) = 255;
            }
        }
    }*/
}

void Classifier::preprocess(const Mat& src, Mat& dst, int i)
{
    Mat img_gray, img_th, img_th2;// img_norm;
    if (src.channels() == 3)
    {
        cv::cvtColor(src, img_gray, COLOR_BGR2GRAY);
    }
    else if (src.channels() == 1)
    {
        img_gray = src;
    }    
    else
    {
        throw "classifier.cpp-->preprocess: src channels error";
    }
    
    int mean = cv::mean(img_gray)[0];
    int thresh = MIN(mean,mean);
    threshold(img_gray, img_th, thresh, 255, THRESH_BINARY);
    //get_max_connectedComponent(img_th, img_th2);
    img_th2 = img_th;
    //normalize(img_th2, img_norm);
    dst = img_th2;

    /* imwrite(path_save2 + to_string(n) + "_" + to_string(save_idx) + ".jpg", src);
     imwrite(path_save2 + to_string(n) + "_" + to_string(save_idx) + "_bin1.jpg", img_th);
     imwrite(path_save2 + to_string(n) + "_" + to_string(save_idx) + "_bin2.jpg", img_th2);*/

    if(is_show)
    {
         Mat show = Mat(src.rows, src.cols *4, CV_8UC3);

        cv::cvtColor(img_gray, img_gray, COLOR_GRAY2BGR);     
        cv::cvtColor(img_th, img_th, COLOR_GRAY2BGR);
        cv::cvtColor(img_th2, img_th2, COLOR_GRAY2BGR);

        int col_idx = 0;
        src.copyTo(show(Rect(src.cols * (col_idx++), 0, src.cols, src.rows)));
        img_gray.copyTo(show(Rect(src.cols * (col_idx++), 0, src.cols, src.rows)));    
        img_th.copyTo(show(Rect(src.cols * (col_idx++), 0, src.cols, src.rows)));
        img_th2.copyTo(show(Rect(src.cols * (col_idx++), 0, src.cols, src.rows)));

        //cv::resize(show, show, Size(0, 0), 2, 2);
        //imshow("shw0000", show);
        //waitKey(1);
        {
            unique_lock<mutex>lock(show_mutex);
            shows[i] = show;
            total_show_rows += show.rows;
            total_show_cols = MAX(show.cols, total_show_cols);
            //cout << "save" <<i<<endl;
        }
    }
}

vector<int> Classifier::operator()(const vector<Mat>& images)
{
    auto procOneFrame = [this,&images](int i)
    {
        //cout << "proc n " << i << endl;
        double t1 = getTickCount();
        Mat img_prep;
        preprocess(images[i], img_prep, i);
        int n = model(img_prep);
        double t2 = getTickCount();
        auto dt = (t2 - t1) / getTickFrequency() * 1000;
        //LOG("NNtime is " + to_string(dt) + "ms\n");
        //std::cout << "NN time: " << dt << " ms" << endl;
        return n;
    };

    if (is_show)
    {
        total_show_cols = 0;
        total_show_rows = 0;
        shows.resize(images.size());
    }

    vector<future<int>>rst_future;
    vector<int> result(images.size());
    if (images.empty())
    {
        return result;
    }

    double t11 = getTickCount();

    
    for (int idx = 1; idx < images.size(); ++idx)
    {       
        //cout << "idx_ " << idx << endl;
        rst_future.emplace_back(threadpool.commit([idx, &procOneFrame] {
            /*cout << "proc n1 " << idx << endl;*/ 
            return procOneFrame(idx);
            }));
    }
    result[0] = procOneFrame(0);
    for (int idx = 0; idx < images.size();++idx)
    {
        if (idx > 0)
        {
            result[idx] = rst_future[idx-1].get();          
        }
#ifdef CLASSIFIER_SHOW
        //Mat show = shows[idx];
        //imshow("show_num", show);
        //imwrite(path_save + to_string(save_idx++) + "_" + to_string(n) + ".jpg", show);
        //waitKey(1); 
#endif
    }
    double t12 = getTickCount();
    auto dt = (t12 - t11) / getTickFrequency() * 1000;
    print("input images: %d , total NNtime : %f ms\n", images.size(), dt);

 /*   double t11 = getTickCount();
    int idx = 0;
    vector<Mat>input(images.size());
    int size = images.size();
    for (int idx = 0; idx < size; ++idx)
    {
        Mat img_th;
        preprocess(images[idx], img_th, idx);
        input[idx] = img_th;
    }
    vector<int> result = model(input);
    double t12 = getTickCount();
    auto dt = (t12 - t11) / getTickFrequency() * 1000;
    LOG("input "+to_string(images.size()) +" images , total NNtime is " + to_string(dt) + "ms\n");
    */

    if (is_show)
    {
        Mat total_show = Mat::zeros(total_show_rows, total_show_cols, CV_8UC3);
        int rows = 0;
        for (int i = 0, size = images.size(); i < size; i++)
        {
            Mat& show = shows[i];
            Point org = {show.cols - 30, 30};
            putText(show, to_string(result[i]), org, FONT_HERSHEY_COMPLEX, 1, Scalar(0, 255, 0));
            show.copyTo(total_show(Rect(0, rows, show.cols, show.rows)));
            rows += show.rows;
        }
        display("shows", total_show);
    }

    return result;
}
